package com.bigdata.flink.statemgmt;

import com.bigdata.flink.util.StreamEnvUtil;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.List;

/**
 *  状态后端 hashMap和rocksDB 默认hashMap内存存储
 */
public class StateMgmtTest {
    public static void main(String[] args) throws Exception {
        reduceStateTest();
    }
    
    // 流可以进行connect，然后共享某些状态
    public static void broadcastState() throws Exception {
        StreamExecutionEnvironment env = StreamEnvUtil.getEnv();
        DataStreamSource<String> socketSource1 = StreamEnvUtil.getSocketSource(env);

        StreamExecutionEnvironment env2 = StreamEnvUtil.getEnv();
        DataStreamSource<String> socketSource2 = StreamEnvUtil.getSocketSource(env2);

        MapStateDescriptor<String, String> state = new MapStateDescriptor<>("state", String.class, String.class);
        BroadcastStream<String> broadcast = socketSource1.broadcast(state);

        SingleOutputStreamOperator<String> process = socketSource1.connect(broadcast)
                .process(new BroadcastProcessFunction<String, String, String>() {
                    @Override
                    public void processElement(String value, ReadOnlyContext ctx, Collector<String> out) throws Exception {
                        BroadcastState<String, String> broadcastState = (BroadcastState<String, String>) ctx.getBroadcastState(state);
                        String max = broadcastState.get("max");

                    }

                    @Override
                    public void processBroadcastElement(String value, Context ctx, Collector<String> out) throws Exception {
                        // 下面这个流可以动态更改配置，影响上面的处理方法，两个方法共享MapState
                        BroadcastState<String, String> broadcastState = ctx.getBroadcastState(state);
                        broadcastState.put("max", "10");
                    }
                });

        process.print();
        env.execute();
    }

    static class CustomState implements MapFunction<String, String>, CheckpointedFunction {

        // 本地变量
        String value;
        // 自定义算子状态
        ListState<String> state;

        @Override
        public String map(String value) throws Exception {
            return value.split("\\s")[0];
        }

        /**
         * 自定义算子状态  checkpoint调用 状态备份
         *
         * @param context the context for drawing a snapshot of the operator
         * @throws Exception
         */
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            // checkpoint把本地变量拷贝到state中等待备份
            state.clear();
            state.add(value);
        }

        /**
         * 自定义算子状态 checkpoint恢复 状态数据
         *
         * @param context the context for initializing the operator
         * @throws Exception
         */
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            // checkpoint恢复状态，并拷贝到本地变量使用
            state = context.getOperatorStateStore()
                    .getListState(new ListStateDescriptor<String>("state", Types.STRING));
            // 随便写的大概就是这个意思
            for (String s : state.get()) {
                value = s;
                break;
            }
        }
    }

    public static void valueStateTest() throws Exception {
        StreamExecutionEnvironment env = StreamEnvUtil.getEnv();
        DataStreamSource<String> socketSource = StreamEnvUtil.getSocketSource(env);

        SingleOutputStreamOperator<Integer> process = socketSource.keyBy(f -> f.split(",")[0])
                .process(new KeyedProcessFunction<String, String, Integer>() {

                    ValueState<Integer> state;

                    @Override
                    public void open(Configuration ctx) throws Exception {
                        super.open(ctx);
                        StateTtlConfig build = StateTtlConfig.newBuilder(Time.seconds(5))
                                .neverReturnExpired()
                                .updateTtlOnCreateAndWrite().build();

                        ValueStateDescriptor<Integer> stateDescriptor = new ValueStateDescriptor<>("state", Types.INT);
                        stateDescriptor.enableTimeToLive(build);
                        state = getRuntimeContext().getState(stateDescriptor);
                    }

                    @Override
                    public void processElement(String key, Context context, Collector<Integer> collector) throws Exception {
                        Integer value = state.value();
                        // 看下先写入 5sTTL后，是否过期，过期这里NULL
                        // 输入格式 s1 1
                        System.out.printf("key: %s, value: %d\n", key, value);
                        int parseInt = Integer.parseInt(key.split("\\s")[1]);
                        state.update(parseInt);
                    }
                });

        process.print();
        env.execute();
    }

    public static void listValueTest() throws Exception {
        StreamExecutionEnvironment env = StreamEnvUtil.getEnv();
        DataStreamSource<String> socketSource = StreamEnvUtil.getSocketSource(env);

        SingleOutputStreamOperator<Integer> process = socketSource.keyBy(f -> f.split("\\s")[0])
                .process(new KeyedProcessFunction<String, String, Integer>() {

                    ListState<Integer> state;

                    @Override
                    public void open(Configuration ctx) throws Exception {
                        super.open(ctx);
                        state = getRuntimeContext().getListState(new ListStateDescriptor<>("state", Types.INT));
                    }

                    @Override
                    public void processElement(String key, Context context, Collector<Integer> collector) throws Exception {
                        // 放一个总和进去，写着玩
                        Iterable<Integer> integers = state.get();

                        List<Integer> list = new ArrayList<>();

                        long size = integers.spliterator().estimateSize();
                        int sum = 0;
                        if (size > 0) {
                            for (Integer integer : integers) {
                                sum += integer;
                            }
                        }

                        String s = key.split("\\s")[1];
                        sum += Integer.parseInt(s);

                        list.add(sum);
                        state.update(list);

                        // 输入格式，如果出现 10, 就输出总和
                        if ("10".equals(s)) {
                            for (Integer integer : state.get()) {
                                collector.collect(integer);
                            }
                        }
                    }
                });

        process.print();
        env.execute();
    }

    public static void mapStateTest() throws Exception {
        StreamExecutionEnvironment env = StreamEnvUtil.getEnv();
        DataStreamSource<String> socketSource = StreamEnvUtil.getSocketSource(env);

        SingleOutputStreamOperator<String> process = socketSource.keyBy(f -> f.split("\\s")[0])
                .process(new KeyedProcessFunction<String, String, String>() {

                    MapState<String, Integer> state;

                    @Override
                    public void open(Configuration ctx) throws Exception {
                        super.open(ctx);
                        state = getRuntimeContext().getMapState(new MapStateDescriptor<>("state", Types.STRING, Types.INT));
                    }

                    @Override
                    public void processElement(String value, Context context, Collector<String> collector) throws Exception {
                        // 统计单词数量, 但某个先到3个，就输出并清空
                        String[] split = value.split("\\s");
                        for (String key : split) {
                            if (state.contains(key)) {
                                Integer count = state.get(key);
                                count++;
                                if (count == 3) {
                                    collector.collect(String.format("key: %s, value: %d", key, count));
                                    state.remove(key);
                                } else {
                                    state.put(key, count);
                                }
                            } else {
                                state.put(key, 1);
                            }
                        }
                    }
                });

        process.print();
        env.execute();
    }

    public static void reduceStateTest() throws Exception {
        StreamExecutionEnvironment env = StreamEnvUtil.getEnv();
        DataStreamSource<String> socketSource = StreamEnvUtil.getSocketSource(env);

        SingleOutputStreamOperator<String> process = socketSource.keyBy(f -> f.split("\\s")[0])
                .process(new KeyedProcessFunction<String, String, String>() {

                    ReducingState<Integer> state;

                    @Override
                    public void open(Configuration ctx) throws Exception {
                        super.open(ctx);
                        state = getRuntimeContext().getReducingState(new ReducingStateDescriptor<>("state",
                                new ReduceFunction<Integer>() {
                                    @Override
                                    public Integer reduce(Integer value1, Integer value2) throws Exception {
                                        // 同样传入第二个参数开始计算
                                        System.out.println("reduce start");
                                        return value1 + value2;
                                    }
                                }, Types.INT));
                    }

                    @Override
                    public void processElement(String value, Context context, Collector<String> collector) throws Exception {
                        // 传值到reduceFunction进行计算
                        String s = value.split("\\s")[1];
                        int parseInt = Integer.parseInt(s);

                        state.add(parseInt);
                        if (parseInt == 10) {
                            collector.collect(String.format("key: %s, value: %d", value.split("\\s")[0], state.get()));
                        }
                    }
                });

        process.print();
        env.execute();
    }

    public static void aggTest() throws Exception {
        StreamExecutionEnvironment env = StreamEnvUtil.getEnv();
        DataStreamSource<String> socketSource = StreamEnvUtil.getSocketSource(env);

        SingleOutputStreamOperator<String> process = socketSource.keyBy(f -> f.split("\\s")[0])
                .process(new KeyedProcessFunction<String, String, String>() {

                    AggregatingState<String, Integer> state;

                    @Override
                    public void open(Configuration ctx) throws Exception {
                        super.open(ctx);
                        state = getRuntimeContext().getAggregatingState(new AggregatingStateDescriptor<String, Integer, Integer>("state", new AggregateFunction<String, Integer, Integer>() {
                            @Override
                            public Integer createAccumulator() {
                                return 0;
                            }

                            @Override
                            public Integer add(String value, Integer accumulator) {
                                return 0;
                            }

                            @Override
                            public Integer getResult(Integer accumulator) {
                                return 0;
                            }

                            @Override
                            public Integer merge(Integer a, Integer b) {
                                return 0;
                            }
                        },  Types.INT));
                    }

                    @Override
                    public void processElement(String value, Context context, Collector<String> collector) throws Exception {
                        // 累加器 略
                    }
                });

        process.print();
        env.execute();
    }
}
